#!/usr/bin/python3
"""
OpenCloudOS 9 系统环境检测工具
用于检测系统是否适合运行AI软件包验证
"""

import subprocess
import sys
import platform
import psutil
import json
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any


class SystemChecker:
    """系统环境检测器"""
    
    def __init__(self):
        self.check_results = {}
        
    def check_os_info(self) -> Dict[str, Any]:
        """检测操作系统信息"""
        print("🔍 检测操作系统信息...")
        
        os_info = {
            "system": platform.system(),
            "release": platform.release(),
            "version": platform.version(),
            "machine": platform.machine(),
            "processor": platform.processor(),
            "python_version": platform.python_version(),
            "check_status": "success"
        }
        
        # 尝试获取发行版信息
        try:
            with open('/etc/os-release', 'r') as f:
                release_info = {}
                for line in f:
                    if '=' in line:
                        key, value = line.strip().split('=', 1)
                        release_info[key] = value.strip('"')
                os_info["distribution"] = release_info
        except:
            os_info["distribution"] = "Unknown"
        
        # 检查是否为OpenCloudOS
        is_opencloudos = False
        if isinstance(os_info.get("distribution"), dict):
            name = os_info["distribution"].get("NAME", "").lower()
            id_like = os_info["distribution"].get("ID_LIKE", "").lower()
            if "opencloudos" in name or "opencloudos" in id_like:
                is_opencloudos = True
        
        os_info["is_opencloudos"] = is_opencloudos
        
        print(f"  ✅ 系统: {os_info['system']} {os_info['release']}")
        print(f"  ✅ Python: {os_info['python_version']}")
        if is_opencloudos:
            print(f"  ✅ 确认为 OpenCloudOS 系统")
        else:
            print(f"  ⚠️  非 OpenCloudOS 系统")
            
        return os_info
    
    def check_hardware(self) -> Dict[str, Any]:
        """检测硬件信息"""
        print("🔍 检测硬件信息...")
        
        hardware_info = {
            "cpu_count": psutil.cpu_count(),
            "cpu_count_logical": psutil.cpu_count(logical=True),
            "memory_total_gb": round(psutil.virtual_memory().total / (1024**3), 2),
            "memory_available_gb": round(psutil.virtual_memory().available / (1024**3), 2),
            "disk_free_gb": round(psutil.disk_usage('/').free / (1024**3), 2),
            "check_status": "success"
        }
        
        print(f"  ✅ CPU核心数: {hardware_info['cpu_count']} 物理 / {hardware_info['cpu_count_logical']} 逻辑")
        print(f"  ✅ 内存: {hardware_info['memory_available_gb']:.1f}GB 可用 / {hardware_info['memory_total_gb']:.1f}GB 总计")
        print(f"  ✅ 磁盘空间: {hardware_info['disk_free_gb']:.1f}GB 可用")
        
        # 检查内存是否足够
        if hardware_info['memory_total_gb'] < 4:
            hardware_info["memory_warning"] = "内存不足4GB，可能影响大型AI包的安装和测试"
            print(f"  ⚠️  {hardware_info['memory_warning']}")
        
        # 检查磁盘空间
        if hardware_info['disk_free_gb'] < 10:
            hardware_info["disk_warning"] = "磁盘空间不足10GB，可能影响AI包的安装"
            print(f"  ⚠️  {hardware_info['disk_warning']}")
        
        return hardware_info
    
    def check_gpu(self) -> Dict[str, Any]:
        """检测GPU信息"""
        print("🔍 检测GPU信息...")
        
        gpu_info = {
            "nvidia_driver_available": False,
            "cuda_available": False,
            "gpu_devices": [],
            "check_status": "success"
        }
        
        # 检查NVIDIA驱动
        try:
            result = subprocess.run(['nvidia-smi'], capture_output=True, text=True)
            if result.returncode == 0:
                gpu_info["nvidia_driver_available"] = True
                print("  ✅ NVIDIA驱动可用")
            else:
                print("  ❌ NVIDIA驱动不可用")
        except FileNotFoundError:
            print("  ❌ 未找到nvidia-smi命令")
        
        # 检查CUDA
        try:
            result = subprocess.run(['nvcc', '--version'], capture_output=True, text=True)
            if result.returncode == 0:
                gpu_info["cuda_available"] = True
                print("  ✅ CUDA工具包可用")
            else:
                print("  ❌ CUDA工具包不可用")
        except FileNotFoundError:
            print("  ❌ 未找到nvcc命令")
        
        return gpu_info
    
    def check_python_env(self) -> Dict[str, Any]:
        """检测Python环境"""
        print("🔍 检测Python环境...")
        
        python_info = {
            "python_version": sys.version,
            "pip_available": False,
            "pip_version": None,
            "virtual_env": None,
            "check_status": "success"
        }
        
        # 检查pip
        try:
            result = subprocess.run([sys.executable, '-m', 'pip', '--version'], 
                                  capture_output=True, text=True)
            if result.returncode == 0:
                python_info["pip_available"] = True
                python_info["pip_version"] = result.stdout.strip()
                print("  ✅ pip可用")
            else:
                print("  ❌ pip不可用")
        except Exception:
            print("  ❌ pip检测失败")
        
        # 检查虚拟环境
        if hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
            python_info["virtual_env"] = "active"
            print("  ✅ 在虚拟环境中")
        else:
            python_info["virtual_env"] = "none"
            print("  ⚠️  未使用虚拟环境")
        
        return python_info
    
    def check_network(self) -> Dict[str, Any]:
        """检测网络连接"""
        print("🔍 检测网络连接...")
        
        network_info = {
            "pypi_accessible": False,
            "internet_available": False,
            "check_status": "success"
        }
        
        # 检查基本网络连接
        try:
            import urllib.request
            urllib.request.urlopen('http://www.baidu.com', timeout=5)
            network_info["internet_available"] = True
            print("  ✅ 互联网连接正常")
        except:
            print("  ❌ 互联网连接异常")
        
        # 检查PyPI连接
        try:
            import urllib.request
            urllib.request.urlopen('https://pypi.org', timeout=10)
            network_info["pypi_accessible"] = True
            print("  ✅ PyPI连接正常")
        except:
            print("  ❌ PyPI连接异常")
        
        return network_info
    
    def run_full_check(self) -> Dict[str, Any]:
        """运行完整系统检测"""
        print("🚀 开始系统环境检测...\n")
        
        self.check_results = {
            "timestamp": datetime.now().isoformat(),
            "os_info": self.check_os_info(),
            "hardware_info": self.check_hardware(),
            "gpu_info": self.check_gpu(),
            "python_info": self.check_python_env(),
            "network_info": self.check_network()
        }
        
        print("\n" + "="*60)
        print("📊 系统检测摘要")
        print("="*60)
        
        # 生成建议
        recommendations = self._generate_recommendations()
        self.check_results["recommendations"] = recommendations
        
        # 保存检测结果
        self._save_results()
        
        return self.check_results
    
    def _generate_recommendations(self) -> List[str]:
        """生成系统优化建议"""
        recommendations = []
        
        # 操作系统建议
        if not self.check_results["os_info"].get("is_opencloudos", False):
            recommendations.append("建议在OpenCloudOS 9系统上运行以获得最佳兼容性")
        
        # 硬件建议
        if self.check_results["hardware_info"]["memory_total_gb"] < 8:
            recommendations.append("建议使用8GB以上内存以支持大型AI模型")
        
        # GPU建议
        if not self.check_results["gpu_info"]["nvidia_driver_available"]:
            recommendations.append("建议安装NVIDIA驱动以支持GPU加速的AI包")
        
        # Python环境建议
        if self.check_results["python_info"]["virtual_env"] == "none":
            recommendations.append("建议使用Python虚拟环境以避免包冲突")
        
        # 网络建议
        if not self.check_results["network_info"]["pypi_accessible"]:
            recommendations.append("建议检查网络连接或配置PyPI镜像源")
        
        if not recommendations:
            recommendations.append("系统环境良好，可以开始AI软件包验证")
        
        for i, rec in enumerate(recommendations, 1):
            print(f"{i}. {rec}")
        
        return recommendations
    
    def _save_results(self):
        """保存检测结果"""
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f"system_check_{timestamp}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(self.check_results, f, ensure_ascii=False, indent=2)
        
        print(f"\n💾 检测结果已保存到: {filename}")


def main():
    """主函数"""
    checker = SystemChecker()
    results = checker.run_full_check()
    
    print("\n🎯 系统检测完成！")
    print("现在可以运行AI软件包验证工具了。")


if __name__ == "__main__":
    main()